package mj

import (
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/imroc/req/v3"
	"gorm.io/gorm"
	"io"
	"myai/core/types"
	logger2 "myai/logger"
	"myai/store/model"
	"myai/utils"
	"time"
)

// Client MidJourney client
type Client struct {
	client *req.Client
	db     *gorm.DB
}

type ImageReq struct {
	// MID_JOURNEY、NIJI_JOURNEY
	BotType string `json:"botType,omitempty"`
	// 提示词
	Prompt string `json:"prompt,omitempty"`
	// 比例: PORTRAIT(2:3); SQUARE(1:1); LANDSCAPE(3:2)
	// 枚举值: PORTRAIT、SQUARE、LANDSCAPE
	Dimensions string `json:"dimensions,omitempty"`
	// 图片base64数组
	Base64Array []string `json:"base64Array,omitempty"`
	// 账号筛选条件
	AccountFilter interface{} `json:"accountFilter,omitempty"`
	// 回调地址, 为空时使用全局notifyHook
	NotifyHook string `json:"notifyHook,omitempty"`
	// 自定义参数
	State string `json:"state,omitempty"`
}

type ImageRes struct {
	Code        int      `json:"code"` // 响应状态码
	Description string   `json:"description"`
	Properties  struct{} `json:"properties"`
	Result      string   `json:"result"`
	Channel     string   `json:"channel,omitempty"` // 请求渠道地址
}

// QueryRes 查询响应
type QueryRes struct {
	// 任务类型 IMAGINE、UPSCALE、VARIATION、ZOOM、PAN、DESCRIBE、BLEND、SHORTEN、SWAP_FACE
	Action string `json:"action"`
	// 可执行按钮
	Buttons []struct {
		CustomId string `json:"customId"` // 动作标识
		Emoji    string `json:"emoji"`    // 图标
		Label    string `json:"label"`    // 文本
		Style    int    `json:"style"`    // 样式
		Type     int    `json:"type"`     // 类型，系统内部使用
	} `json:"buttons"`
	Description string   `json:"description"` // 任务描述
	FailReason  string   `json:"failReason"`  // 失败原因
	FinishTime  int      `json:"finishTime"`  // 结束时间
	Id          string   `json:"id"`          // Id
	ImageUrl    string   `json:"imageUrl"`    // 图片URL
	Progress    string   `json:"progress"`    // 任务进度
	Prompt      string   `json:"prompt"`      // 提示词
	PromptEn    string   `json:"promptEn"`    // 提示词英文
	Properties  struct{} `json:"properties"`
	StartTime   int      `json:"startTime"`  // 开始执行时间
	State       string   `json:"state"`      // 自定义参数
	Status      string   `json:"status"`     // 任务状态 NOT_START、SUBMITTED、MODAL、IN_PROGRESS、FAILURE、SUCCESS、CANCEL
	SubmitTime  int      `json:"submitTime"` // 提交时间
}

var log = logger2.GetLogger()

func NewClient(db *gorm.DB) *Client {
	return &Client{
		client: req.C().SetTimeout(time.Minute).SetUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"),
		db:     db,
	}
}

// Imagine 绘图
func (c *Client) Imagine(task types.MjTask) (ImageRes, error) {
	apiPath := fmt.Sprintf("mj-%s/mj/submit/imagine", task.Mode)
	prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
	// 设置负面提示词
	if task.NegPrompt != "" {
		prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
	}
	body := ImageReq{
		BotType:     "MID_JOURNEY",
		Prompt:      prompt,
		Base64Array: make([]string, 0), // 垫图base64数组
	}
	// 生成图片 Base64 编码
	if len(task.ImgArr) > 0 {
		imageData, err := utils.DownloadImg(task.ImgArr[0], "")
		if err != nil {
			log.Errorf("下载图片出错：%v", err)
		} else {
			body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
		}
	}
	return c.doRequest(body, apiPath, task.ChannelId)
}

// Blend 融图
func (c *Client) Blend(task types.MjTask) (ImageRes, error) {
	apiPath := fmt.Sprintf("mj-%s/mj/submit/blend", task.Mode)
	body := ImageReq{
		// MID_JOURNEY、NIJI_JOURNEY
		BotType: "MID_JOURNEY",
		// 比例: PORTRAIT(2:3); SQUARE(1:1); LANDSCAPE(3:2)
		// 枚举值: PORTRAIT、SQUARE、LANDSCAPE
		Dimensions: "SQUARE",
		// 图片base64数组
		Base64Array: make([]string, 0),
	}
	// 生成图片Base64编码
	if len(task.ImgArr) > 0 {
		for _, imgURL := range task.ImgArr {
			// 下载图片
			imageData, err := utils.DownloadImg(imgURL, "")
			if err != nil {
				log.Errorf("图片下载失败：%v", err)
			} else {
				body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
			}
		}
	}
	return c.doRequest(body, apiPath, task.ChannelId)
}

// SwapFace 换脸
func (c *Client) SwapFace(task types.MjTask) (ImageRes, error) {
	apiPath := fmt.Sprintf("mj-%s/mj/insight-face/swap", task.Mode)
	// 生成图片 Base64 编码
	if len(task.ImgArr) != 2 {
		return ImageRes{}, errors.New("参数错误，必须上传2张图片才能进行换脸")
	}
	var sourceBase64 string
	var targetBase64 string
	// 下载图片
	imageData, err := utils.DownloadImg(task.ImgArr[0], "")
	if err != nil {
		log.Errorf("图片下载失败：%v", err)
	} else {
		// 把图片转成base64格式
		sourceBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
	}
	// 下载图片
	imageData, err = utils.DownloadImg(task.ImgArr[1], "")
	if err != nil {
		log.Errorf("图片下载失败：%v", err)
	} else {
		targetBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
	}

	body := gin.H{
		"sourceBase64": sourceBase64, // 人脸源图片base64
		"targetBase64": targetBase64, // 目标图片base64
		"accountFilter": gin.H{
			"instanceId": "", // 账号实例ID
		},
		"state": "", // 自定义参数
	}
	return c.doRequest(body, apiPath, task.ChannelId)
}

// Upscale 放大指定的图片
func (c *Client) Upscale(task types.MjTask) (ImageRes, error) {
	body := map[string]string{
		"customId": fmt.Sprintf("MJ::JOB::upsample::%d::%s", task.Index, task.MessageHash),
		"taskId":   task.MessageId,
	}
	apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
	return c.doRequest(body, apiPath, task.ChannelId)
}

// Variation 以指定的图片的视角进行变换在创作，注意主要在对应的频道中关闭 Remix 变换，否则 Variation 执行不会生效
func (c *Client) Variation(task types.MjTask) (ImageRes, error) {
	body := map[string]string{
		"customId": fmt.Sprintf("MJ::JOB::variation::%d::%s", task.Index, task.MessageHash),
		"taskId":   task.MessageId,
	}
	apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
	return c.doRequest(body, apiPath, task.ChannelId)
}

// doRequest 发送请求
func (c *Client) doRequest(body interface{}, apiPath string, channel string) (ImageRes, error) {
	var res ImageRes
	session := c.db.Session(&gorm.Session{}).Where("type", "mj").Where("enabled", true)
	if channel != "" {
		session = session.Where("api_url", channel)
	}
	var apiKey model.ApiKey
	err := session.Order("last_used_at ASC").First(&apiKey).Error
	if err != nil {
		return ImageRes{}, fmt.Errorf("没有可用的MidJourney api密钥：%v", err)
	}

	apiURL := fmt.Sprintf("%s/%s", apiKey.ApiURL, apiPath)
	log.Infof("MJ请求地址：%s", apiURL)
	r, err := req.C().R().SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(body).
		SetSuccessResult(&res).
		Post(apiURL)
	if err != nil {
		return ImageRes{}, fmt.Errorf("请求API:%s 出错：%v", apiURL, err)
	}
	if r.IsErrorState() {
		// 读取响应错误消息
		errMsg, _ := io.ReadAll(r.Body)
		return ImageRes{}, fmt.Errorf("API 返回错误：%s", string(errMsg))
	}
	// 更新API KEY最后一次使用的时间
	if err = c.db.Model(&apiKey).Update("last_used_at", time.Now().Unix()).Error; err != nil {
		log.Errorf("更新API KEY最后一次使用的时间出错：%v", err)
	}
	res.Channel = apiKey.ApiURL
	return res, nil
}

// QueryTask 查询指定ID任务
func (c *Client) QueryTask(taskId string, channel string) (QueryRes, error) {
	var apiKey model.ApiKey
	err := c.db.Where("type", "mj").Where("enabled", true).Where("api_url", channel).First(&apiKey).Error
	if err != nil {
		return QueryRes{}, fmt.Errorf("没有可用的MidJourney api密钥：%v", err)
	}
	apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", apiKey.ApiURL, taskId)
	var res QueryRes
	r, err := c.client.R().SetHeader("Authorization", "Bearer "+apiKey.Value).SetSuccessResult(&res).Get(apiURL)
	if err != nil {
		return QueryRes{}, fmt.Errorf("查询任务请求失败：%v", err)
	}
	if r.IsErrorState() {
		return QueryRes{}, errors.New("error status:" + r.Status)
	}
	return res, nil
}
